package com.caoyanan.algorithm.question.zuoTraining.training003.class06;

import com.caoyanan.algorithm.question.zuoTraining.LogarithmInputGenerator;

import java.util.Arrays;
import java.util.Comparator;
import java.util.PriorityQueue;

/**
 * 给两个有序数组，再给一个整数K。
 * 求两个数累加和最大的前K个，两个数必须分别来自arr1和arr2
 *
 * 比如 arr1 = [1,5,9,12]， arr2 = [3,6,7,14]， k = 3
 * 累加和排第一的是 arr1[3] + arr2[3] = 12 + 14 = 26
 * 排在第二的是  arr1[9] + arr2[3] = 9 + 12 = 21
 * 即arr2[3]可以重复用，但是组合不可以重复用
 *
 * 一个样本做行，一个样本做列的对应模型
 *
 * @author: caoyanan
 * @time: 2021/4/25 3:23 下午
 */
public class Question05_TopKSumCrossTwoArrays {

    public static class Node {
        private final int rowIndex;
        private final int columnIndex;
        private final int sum;

        public Node(int rowIndex, int columnIndex, int sum) {
            this.rowIndex = rowIndex;
            this.columnIndex = columnIndex;
            this.sum = sum;
        }

        public int getSum() {
            return sum;
        }
    }

    public static class MaxHeapComparator implements Comparator<Node> {

        @Override
        public int compare(Node o1, Node o2) {
            return o2.sum - o1.sum;
        }
    }


    /**
     * 大顶堆 O(k * log(k))
     * @return topK
     */
    public static int[] topkSumCrossTwoArray(int[] arr1, int[] arr2, int k) {

        if (arr1.length == 0 || arr2.length == 0 || k == 0) {
            return new int[0];
        }
        k = Math.min(k, arr1.length * arr2.length);
        PriorityQueue<Node> nodes = new PriorityQueue<>(new MaxHeapComparator());
        int arr1MaxIndex = arr1.length - 1;
        int arr2MaxIndex = arr2.length - 1;
        nodes.add(new Node(arr1MaxIndex, arr2MaxIndex, arr1[arr1MaxIndex] + arr2[arr2MaxIndex]));

        int[] result = new int[k];
        int resIndex = 0;
        boolean[][] visited = new boolean[arr1.length][arr2.length];
        while (k != resIndex && nodes.size() != 0) {
            Node maxSum = nodes.poll();
            result[resIndex++] = maxSum.getSum();
            if (maxSum.rowIndex > 0 && !visited[maxSum.rowIndex - 1][maxSum.columnIndex]) {
                nodes.add(new Node(maxSum.rowIndex - 1, maxSum.columnIndex,
                        arr1[maxSum.rowIndex - 1] + arr2[maxSum.columnIndex]));
                visited[maxSum.rowIndex - 1][maxSum.columnIndex] = true;
            }
            if (maxSum.columnIndex > 0 && !visited[maxSum.rowIndex][maxSum.columnIndex - 1]) {
                nodes.add(new Node(maxSum.rowIndex, maxSum.columnIndex - 1,
                        arr1[maxSum.rowIndex] + arr2[maxSum.columnIndex - 1]));
                visited[maxSum.rowIndex][maxSum.columnIndex - 1] = true;
            }
        }
        return result;
    }

    /**
     * 暴力法 O(n^2)
     * @param arr1
     * @param arr2
     * @param k
     * @return
     */
    public static int[] forTest(int[] arr1, int[] arr2, int k) {
        k = Math.min(k, arr1.length * arr2.length);
        int[] allSum = new int[arr1.length * arr2.length];
        int index = 0;
        for (int i : arr1) {
            for (int j : arr2) {
                allSum[index++] = i + j;
            }
        }
        Arrays.sort(allSum);
        int[] res = new int[Math.min(k, allSum.length)];
        index = allSum.length - 1;
        for (int i = 0; i != res.length; i++) {
            res[i] = allSum[index--];
        }
        return res;
    }

    public static void main(String[] args) {
        LogarithmInputGenerator instance = LogarithmInputGenerator.getInstance();
        int k = 5;
        for (int i = 0; i < 20; i++) {
            int[] arr1 = instance.generateRandomSortArray(i, 100);
            int[] arr2 = instance.generateRandomSortArray(i, 100);
            int[] absoluteCorrectResult = forTest(arr1, arr2, k);
            int[] answer = topkSumCrossTwoArray(arr1, arr2, k);
            System.out.println("arr1:" + Arrays.toString(arr1));
            System.out.println("arr2:" + Arrays.toString(arr2));
            System.out.printf("violence method top %s: %s%n", k, Arrays.toString(absoluteCorrectResult));
            System.out.printf("heap method top %s: %s%n", k, Arrays.toString(answer));
            System.out.println("是否相等: " + instance.isArrayEqual(absoluteCorrectResult, answer));
            System.out.println("\n");
        }

    }


}
